#include "UdpSocketHandle.h"
#include <cstring>
#include <iostream>
#include <stdexcept>
#include "utils.h"

static constexpr size_t ReadBufferSize = 65536;

thread_local static uint8_t ReadBuffer[ReadBufferSize];

// 接收数据缓冲区函数
inline static void onAlloc(uv_handle_t* handle, size_t suggested_size, uv_buf_t* buf) {
    
    auto *socket = static_cast<UdpSocketHandle*>(handle->data);

    if (socket) {
        socket->OnUvRecvAlloc(suggested_size, buf);
    }
}

// 数据接收回调函数
inline static void onRecv(uv_udp_t* handle, ssize_t nread, const uv_buf_t* buf,
                          const struct sockaddr* addr, unsigned flags) {
    
    auto *socket = static_cast<UdpSocketHandle*>(handle->data);

    if (socket) {
        socket->OnUvRecv(nread, buf, addr, flags);
    }
}

inline static void onSend(uv_udp_send_t* req, int status) {
    
    auto* sendData = static_cast<UdpSocketHandle::UvSendData*>(req->data);
    auto *handle = req->handle;

    auto *socket = static_cast<UdpSocketHandle*>(req->data);

    const auto *cb = sendData->cb;

    if (socket) {
        socket->OnUvSend(status, cb);
    }

    delete sendData;
}

inline static void onCloseUdp(uv_handle_t* handle) {
    delete reinterpret_cast<uv_udp_t *>(handle);
}

/* Instance methods. */

UdpSocketHandle::UdpSocketHandle(uv_udp_t *handle):uvHandle(handle) {

    this->uvHandle->data = static_cast<void *>(this);
    
    int err = uv_udp_recv_start(this->uvHandle, static_cast<uv_alloc_cb>(onAlloc), static_cast<uv_udp_recv_cb>(onRecv));

    if (err != 0) {
        uv_close(reinterpret_cast<uv_handle_t *>(this->uvHandle), static_cast<uv_close_cb>(onCloseUdp));
        throw std::runtime_error("Failed to start UDP receive: " + std::string(uv_strerror(err)));
    }

    if (!SetLocalAddress()) {
        uv_close(reinterpret_cast<uv_handle_t *>(this->uvHandle), static_cast<uv_close_cb>(onCloseUdp));
        throw std::runtime_error("Failed to set local address for UDP socket");
    }
}

UdpSocketHandle::~UdpSocketHandle() {
    if (!this->closed) {
        InternalClose();
    }
}


void UdpSocketHandle::Send(const uint8_t* data, size_t len, const struct sockaddr *addr, UdpSocketHandle::onSendCallback *cb) {
    if (this->closed) {
        if (cb) {
            (*cb)(false);
            delete cb;
        }
        return;
    }
    if (len == 0) {
        if (cb) {
            (*cb)(true);
            delete cb;
        }
        return;
    }

    uv_buf_t buffer = uv_buf_init(reinterpret_cast<char *>(const_cast<uint8_t *>(data)), len);
    const int sent = uv_udp_try_send(this->uvHandle, &buffer, 1, addr);

    if (sent == static_cast<int>(len)) {
        this->sentBytes += len;
        if (cb) {
            (*cb)(true);
            delete cb;
        }
        return;
    } else if (sent == 0){
        this->sentBytes += sent;
        if (cb) {
            (*cb)(false);
            delete cb;
        }
        return;
    } else if (sent != UV_EAGAIN) {
        std::cout << "uv_udp_try_send failed: " << uv_strerror(sent) << std::endl;
    }

    auto *sendData = new UdpSocketHandle::UvSendData(len);

    sendData->req.data = static_cast<void *>(sendData);
    std::memcpy(sendData->store, data, len);
    sendData->cb = cb;

    buffer = uv_buf_init(reinterpret_cast<char *>(sendData->store), len);
    const int err = uv_udp_send(&sendData->req, this->uvHandle, &buffer, 1, addr, static_cast<uv_udp_send_cb>(onSend));

    if (err != 0) {
        std::cout << "uv_udp_send failed: " << uv_strerror(err) << std::endl;
        if (cb) {
            (*cb)(false);
            delete cb;
        }
        delete sendData;
    } else {
        this->sentBytes += len;
    }
}

uint32_t UdpSocketHandle::GetSendBufferSize() const {
    int size = 0;

    const int err = uv_send_buffer_size(reinterpret_cast<uv_handle_t *>(this->uvHandle), &size);

    if (err) {
        std::cout << "Failed to get send buffer size" << std::endl;
        return 0;
    }
    return static_cast<uint32_t>(size);
}

void UdpSocketHandle::SetSendBufferSize(uint32_t size) {

    auto sizeInt = static_cast<int>(size);

    if (sizeInt <= 0) {
        std::cout << "Invalid send buffer size: " << size << std::endl;
    }


    int err = uv_send_buffer_size(reinterpret_cast<uv_handle_t *>(this->uvHandle), std::addressof(sizeInt));

    if (err) {
        std::cout << "Failed to set send buffer size: " << uv_strerror(err) << std::endl;
    }
}

uint32_t UdpSocketHandle::GetRecvBufferSize() const {
    int size = 0;

    const int err = uv_recv_buffer_size(reinterpret_cast<uv_handle_t *>(this->uvHandle), std::addressof(size));

    if (err) {
        std::cout << "Failed to get receive buffer size" << std::endl;
    }
    return static_cast<uint32_t>(size);
}

void UdpSocketHandle::SetRecvBufferSize(uint32_t size) {

    auto sizeInt = static_cast<int>(size);

    if (sizeInt <= 0) {
        std::cout << "Invalid receive buffer size: " << size << std::endl;
    }

    int err = uv_recv_buffer_size(reinterpret_cast<uv_handle_t *>(this->uvHandle), std::addressof(sizeInt));

    if (err) {
        std::cout << "Failed to set receive buffer size: " << uv_strerror(err) << std::endl;
    }
}

void UdpSocketHandle::InternalClose() {
    if (this->closed) {
        return;
    }

    this->closed = true;

    this->uvHandle->data = nullptr;



    const int err = uv_udp_recv_stop(this->uvHandle);

    if (err != 0) {
        std::cout << "Failed to stop UDP receive: " << uv_strerror(err) << std::endl;
    }

    uv_close(reinterpret_cast<uv_handle_t *>(this->uvHandle), static_cast<uv_close_cb>(onCloseUdp));
}

bool UdpSocketHandle::SetLocalAddress() {

    int err;
    int len = sizeof(this->localAddr);

    err = uv_udp_getsockname(this->uvHandle, reinterpret_cast<struct sockaddr *>(&this->localAddr), &len);

    if (err != 0) {
        std::cout << "Failed to get local address" << std::endl;
        return false;
    }

    int family;

    Utils::IP::GetAddressInfo(reinterpret_cast<const struct sockaddr *>(&this->localAddr), family, this->localIp, this->localPort);

    return true;
}

inline void UdpSocketHandle::OnUvRecvAlloc(size_t /*suggested_size*/, uv_buf_t* buf) {


    buf->base = reinterpret_cast<char *>(ReadBuffer);
    buf->len = ReadBufferSize;
}

inline void UdpSocketHandle::OnUvRecv(ssize_t nread, const uv_buf_t* buf, const struct sockaddr* addr, unsigned flags) {

    if (nread == 0) {
        return;
    }
    if (flags & UV_UDP_PARTIAL != 0u) {
        std::cout << "Received partial UDP datagram" << std::endl;
        return;
    }

    if (nread > 0) {
        this->recvBytes += nread;

        UserOnUdpDatagramReceived(reinterpret_cast<const uint8_t *>(buf->base), static_cast<size_t>(nread), addr);
    } else {
         std::cout << "UDP receive error: " << uv_strerror(nread) << std::endl;
    }
}

inline void UdpSocketHandle::OnUvSend(int status, UdpSocketHandle::onSendCallback* cb) {

    if (status == 0) {
        if (cb) {
            (*cb)(true);
            delete cb;
        }
    } else {
        std::cout << "UDP send error: " << uv_strerror(status) << std::endl;
        if (cb) {
            (*cb)(false);
        }
    }
}